import psycopg2
import json
from psycopg2.extras import Json
from loguru import logger

# 配置日志
logger.add("db_setup.log", rotation="1 MB", level="DEBUG")

# 数据库连接参数
db_params = {
    "database": "water",
    "user": "postgres",
    "password": "water",
    "host": "10.48.0.85",
    "port": "5432"
}

# 要存储的停用词集合
stop_words = [
    "的", "了", "在", "是", "我", "有", "和", "就", "不", "人", "都", "一", "一个", "上", "也", "很", "到",
    "说", "要", "去", "你", "会", "着", "没有", "看", "好", "自己", "这", "那", "什么", "多少", "怎么",
    "为什么", "哪里", "什么时候", "如何", "请", "给我", "提供", "告诉我", "我想知道", "能否", "可以", "吗",
    "呢", "啊", "呀", "吧", "么", "嘛", "及", "以上", "以下", "等", "包括", "主要", "基本", "相关"
]

# 要存储的问题列表
reference_questions = [
    "大伙房水库流域面积多少？",
    "大伙房水库设计标准？",
    "大伙房水库工程规模",
    "大伙房水库总库容/防洪库容/兴利库容/死库容？",
    "大伙房水库校核水位/设计水位/防洪高水位/正常高（兴利）水位/死水位？",
    "大伙房水库主要有哪些建筑物？",
    "大伙房水库汛限水位/防洪限制水位？",
    "编制《大伙房水库调度规程》的主要目的？",
    "《调度规程》适用哪些调度活动？",
    "供水城市有哪些？",
    "调度目标和任务？",
    "主管部门？",
    "调度单位？",
    "运行管理单位？",
    "防洪调度核心任务？",
    "供水调度核心任务？",
    "沈阳断面300年一遇洪水安全泄量？",
    "抚顺断面200年洪水安全泄量？",
    "主汛期防洪限制水位？",
    "前汛期汛限水位？",
    "后汛期调度方式？",
    "超标准洪水调度方式？",
    "安全监测与巡视内容？",
    "库水位监测频次？",
    "降水量观测设备？",
    "自动测报系统通信方式？",
    "短期洪水预报预见期？",
    "DHF预报模型精度？",
    "中长期水文预报精度？",
    "编制《大伙房水库调度规程》的法律依据包括哪些？",
    "《大伙房水库调度规程》技术规范中引用了哪些国家标准或行业导则？",
    "大伙房水库调度的指导思想/原则是什么？",
    "大伙房水库调度的基本原则有哪些？",
    "辽宁省大伙房水库管理局的具体职责是什么？",
    "大伙房水库运行管理单位在汛期的调度权限如何？",
    "大伙房水库主溢洪道弧形工作闸门的孔口尺寸、设计水头、设计水位、堰顶高程是多少？",
    "大伙房水库非常溢洪道的孔口尺寸、设计水头、设计水位、堰顶高程分别是多少？",
    "大伙房水库变形观测的水平位移/垂直位移采用什么方法？",
    "大伙房水库降水和蒸发的观测形式是什么？",
    "大伙房水库水文自动测报系统覆盖的遥测站点有哪些？",
    "大伙房降雨径流模型（DHF）的预报精度如何？",
    "大伙房水库防洪调度时段如何划分？",
    "大伙房水库下游防洪控制断面有哪些？",
    "大伙房水库库水位达到126.4m/128.8m/131.5/133.0/136.32/136.52时如何调度？",
    "大伙房水库自哪一年开始试行分期防洪调度？",
    "大伙房水库非常溢洪道的启用标准是什么？",
    "大伙房水库1000年一遇洪水洪峰流量、7日洪量分别是多少？可能最大洪水洪峰流量、7日洪量分别是多少？",
    "大伙房水库供水调度的根本原则是什么？",
    "大伙房水库城市生活及工业供水保证率是多少？农业灌溉保证率是多少？",
    "大伙房水库生态供水的调度方式是什么？",
    "大伙房水库输水工程调度的原则是什么？",
    "大伙房水库供水对象的优先级顺序是什么？",
    "大伙房水库下游农业灌溉面积和年供水量是多少？",
    "大伙房水库输水工程二期一步的年输水规模是多少？",
    "大伙房水库非汛期生态基流的最低要求是多少？",
    "大伙房水库特殊干旱年的供水调度原则是什么？",
    "大伙房水库抗旱预警水位是多少？触发后采取什么措施？"
]

try:
    # 连接数据库
    conn = psycopg2.connect(**db_params)
    logger.info("成功连接到数据库")

    with conn.cursor() as cursor:
        # 准备停用词数据（使用Json适配器）
        stop_words_data = Json(stop_words)

        # 准备问题数据
        questions_json = json.dumps(reference_questions, ensure_ascii=False)

        # 检查记录是否存在
        cursor.execute("SELECT id FROM t_nlp_config WHERE id = 1")
        record_exists = cursor.fetchone()

        if record_exists:
            # 更新现有记录
            update_query = """
            UPDATE t_nlp_config 
            SET stop_words = %s, 
                reference_questions = %s
            WHERE id = 1
            RETURNING *;
            """
            cursor.execute(update_query, (stop_words_data, questions_json))
            logger.info(f"成功更新了ID为1的记录")
        else:
            # 创建新记录
            insert_query = """
            INSERT INTO t_nlp_config 
            (id, stop_words, reference_questions) 
            VALUES (1, %s, %s) 
            RETURNING *;
            """
            cursor.execute(insert_query, (stop_words_data, questions_json))
            logger.info(f"成功创建了ID为1的新记录")

        # 获取并打印更新后的结果
        result = cursor.fetchone()
        logger.info(f"\n更新后的配置记录:")
        logger.info(f"ID: {result[0]}")
        logger.info(f"停用词数量: {len(result[1])}")
        logger.info(f"参考问题数量: {len(result[2])}")

        logger.info("\n表中的所有记录:")
        for row in cursor.fetchall():
            stop_count = len(row[1]) if row[1] else 0
            ref_count = len(row[2]) if row[2] else 0
            logger.info(f"ID {row[0]}: {stop_count}个停用词, {ref_count}个参考问题")

        # 提交事务
        conn.commit()
        logger.info("事务已提交")

except psycopg2.Error as e:
    logger.error(f"数据库操作出错: {e}")
    if 'conn' in locals() and not conn.closed:
        conn.rollback()
        logger.info("事务已回滚")
    raise
except Exception as e:
    logger.error(f"程序执行出错: {e}")
    raise
finally:
    # 确保连接被关闭
    if 'conn' in locals() and not conn.closed:
        conn.close()
        logger.info("数据库连接已关闭")